import os.path

import onnx
import onnxruntime
from PIL import Image
import numpy as np
import cv2


def post_process(probability, threshold, min_size):
    """
    对每个预测掩码进行后处理时，将忽略像素数小于 'min_size' 的
    :param probability: 预测的mask
    :param threshold: 阈值
    :param min_size:
    :return:
    """
    mask = cv2.threshold(probability, threshold, 1, cv2.THRESH_BINARY)[1]
    num_component, component = cv2.connectedComponents(mask.astype(np.uint8))
    predictions = np.zeros(probability.shape, np.float32)
    num = 0
    for c in range(1, num_component):
        p = (component == c)
        if p.sum() > min_size:
            predictions[p] = 1
            num += 1
    return predictions, num


def stack_mask(image, mask, class_color):
    """
    将预测mask添加到原图
    :param image: b,3,h,w numpy.ndarray
    :param mask: b,num_class,h,w numpy.ndarray
    :param class_color: num_class,3 numpy.ndarray，包含每个类的颜色
    :return: numpy.ndarray
    """
    batch_size = image.shape[0]
    res = []

    for index in range(batch_size):
        img = (image[index]).transpose((1, 2, 0))
        ma = mask[index]

        # 叠加
        mm = np.zeros((ma.shape[1], ma.shape[2], 3), dtype=np.float32)  # shape: (h, w, 3)

        for i in range(ma.shape[0]):
            ma_0 = ma[i] * class_color[i][0]
            ma_1 = ma[i] * class_color[i][1]
            ma_2 = ma[i] * class_color[i][2]

            mm += np.stack([ma_0, ma_1, ma_2], axis=2)

        img_ma = np.clip(img + mm, 0, 1)  # 确保结果在 [0, 1] 之间
        res.append(img_ma.transpose((2, 0, 1)))

    return np.stack(res, axis=0)  # shape: (b, 3, h, w)


class OnnxInference:
    def __init__(self, onnx_model_path, config):
        self.config = config

        self.input_names = ["input"]
        self.output_names = ["output"]
        self.batch_size = 1

        onnx_model = onnx.load(onnx_model_path)
        self.session = onnxruntime.InferenceSession(onnx_model.SerializeToString())

    def normalize_img_fuc(self, img):
        mean = np.array(self.config.dataset.mean)
        std = np.array(self.config.dataset.std)
        img = (img - mean[:, np.newaxis, np.newaxis]) / std[:, np.newaxis, np.newaxis]
        img = img[np.newaxis, :, :, :]
        img = np.array(img, dtype=np.float32)
        return img

    def load_img(self, img_path):
        self.file_name = img_path
        img = np.array(
            Image.open(img_path).convert('RGB').resize((self.config.dataset.width, self.config.dataset.height))) / 255.0
        img = img.transpose((2, 0, 1))
        self.img = img[np.newaxis, ...]
        self.normalize_img = self.normalize_img_fuc(img)
        return self

    def tta_predict(self, seg=True):
        origin_pred = self.session.run(self.output_names, {self.input_names[0]: self.normalize_img})[0]
        preds = np.zeros_like(origin_pred)
        # 水平反转
        img_hlfp = np.flip(self.normalize_img, axis=3)
        hlfp_pred = self.session.run(self.output_names, {self.input_names[0]: img_hlfp})[0]
        # 垂直翻转
        img_vflip = np.flip(self.normalize_img, axis=2)
        vflip_pred = self.session.run(self.output_names, {self.input_names[0]: img_vflip})[0]
        # 水平加垂直翻转

        if seg:
            hlfp_pred = np.flip(hlfp_pred, axis=3)
            vflip_pred = np.flip(vflip_pred, axis=2)

        preds = preds + hlfp_pred + vflip_pred
        self.predict = preds / 3.0
        return self

    def sample_predict(self):
        self.predict = self.session.run(self.output_names, {self.input_names[0]: self.normalize_img})[0]
        return self

    def save_res(self, save_folder):
        pass


class OnnxInferenceClassify(OnnxInference):
    def save_res(self, save_folder):
        preds = self.predict
        class_name = self.config.dataset.class_name
        res = []
        for i in preds:
            if i != 0:
                res.append(class_name[i])
        self.res = res

        with open(os.path.join(save_folder, os.path.basename(self.file_name) + 'classify_txt'), 'w') as f:
            f.write(str(res))


class OnnxInferenceSegment(OnnxInference):

    def save_res(self, save_folder):
        preds = self.predict
        for index, predict_masks_classes in enumerate(preds):
            for each_class, pred in enumerate(predict_masks_classes):
                pred_binary, _ = post_process(pred, self.config.best_thresholds[each_class],
                                              self.config.best_minareas[each_class])
                preds[index, each_class] = pred_binary
        img_pred = stack_mask(self.img, preds, np.array(self.config.dataset.mask_color) / 255.0)
        img_pred = (img_pred[0].transpose((1, 2, 0)) * 255).astype(np.uint8)
        self.res = img_pred
        cv2.imwrite(os.path.join(save_folder, os.path.basename(self.file_name) + '_mask.png'), img_pred)


class OnnxInferenceSegmentClassify(OnnxInferenceSegment):
    def __init__(self, onnx_model_path, additional_onnx_model_path, config):
        assert 'classify_model' in config, "classify_model is not set"
        super().__init__(onnx_model_path, config)

        additional_onnx_model = onnx.load(additional_onnx_model_path)
        self.additional_session = onnxruntime.InferenceSession(additional_onnx_model.SerializeToString())

    def sample_predict(self):
        predict = self.session.run(self.output_names, {self.input_names[0]: self.normalize_img})[0]
        additional_predict = self.additional_session.run(self.output_names, {self.input_names[0]: self.normalize_img})[
            0]
        for i, classes in enumerate(additional_predict):
            for each_class, pred in enumerate(classes):
                if pred == 0:
                    predict[i, each_class, ...] = 0
        self.predict = predict

    def tta_predict(self, seg=True):
        origin_pred = self.session.run(self.output_names, {self.input_names[0]: self.normalize_img})[0]
        preds = np.zeros_like(origin_pred)
        # 水平反转
        img_hlfp = np.flip(self.normalize_img, axis=3)
        hlfp_pred = self.session.run(self.output_names, {self.input_names[0]: img_hlfp})[0]
        # 垂直翻转
        img_vflip = np.flip(self.normalize_img, axis=2)
        vflip_pred = self.session.run(self.output_names, {self.input_names[0]: img_vflip})[0]
        # 水平加垂直翻转
        preds = preds + hlfp_pred + vflip_pred
        predict = preds / 3.0

        origin_pred = self.additional_session.run(self.output_names, {self.input_names[0]: self.normalize_img})[0]
        preds = np.zeros_like(origin_pred)
        # 水平反转
        hlfp_pred = self.additional_session.run(self.output_names, {self.input_names[0]: img_hlfp})[0]
        # 垂直翻转
        vflip_pred = self.additional_session.run(self.output_names, {self.input_names[0]: img_vflip})[0]
        # 水平加垂直翻转

        preds = preds + hlfp_pred + vflip_pred
        additional_predict = preds / 3.0
        for i, classes in enumerate(additional_predict):
            for each_class, pred in enumerate(classes):
                if pred == 0:
                    predict[i, each_class, ...] = 0
        self.predict = predict


def inference_factory(config):
    if config.task_type == 'classify':
        return OnnxInferenceClassify('./export.onnx', config)
    elif config.task_type == 'segment':
        if 'classify_model' in config:
            return OnnxInferenceSegmentClassify('./export.onnx', './additional_export.onnx', config)
        return OnnxInferenceSegment('./export.onnx', config)
    elif config.task_type == 'unsupervised':
        pass


if __name__ == '__main__':
    from omegaconf import OmegaConf

    config = OmegaConf.load("config.yaml")
    app = OnnxInferenceClassify('./export.onnx', config)
    app.load_img('./1e446e941.jpg').tta_predict().save_res('./')
